import os
import sys
import io
import time

import paramiko
import json

# 获取建木参数
# ssh连接参数
ssh_ip = os.getenv("JIANMU_SSH_IP")
ssh_port = os.getenv("JIANMU_SSH_PORT")

try:
    ssh_port = int(ssh_port)
except Exception as e:
    print("ssh端口获取错误：", e, flush=True)
    sys.exit(1)

ssh_user = os.getenv("JIANMU_SSH_USER")
ssh_pass = os.getenv("JIANMU_SSH_PASS")
ssh_private_key = os.getenv("JIANMU_SSH_PRIVATE_KEY")

# 启动参数
jupyter_port = os.getenv("JIANMU_JUPYTER_PORT")
shm_size = os.getenv("JIANMU_SHM_SIZE")
memlock = os.getenv("JIANMU_CONTAINER_MEMLOCK")
stack = os.getenv("JIANMU_CONTAINER_STACK")

nemo_image = "docker.jianmuhub.com/autoops/nemo:latest"
driver_check_cmd = "nvidia-smi"
docker_path = "docker"
container_name = "autoops-nemo"
timeout_seconds = 300


def getSSHClient():
    """获取ssh client"""
    ssh_client = paramiko.SSHClient()
    ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        if ssh_private_key:
            pkey = paramiko.RSAKey.from_private_key(io.StringIO(ssh_private_key))
            ssh_client.connect(hostname=ssh_ip, port=ssh_port, username=ssh_user, pkey=pkey)
            print("Get ssh_client by pkey.", flush=True)
        elif ssh_pass:
            ssh_client.connect(hostname=ssh_ip, port=ssh_port, username=ssh_user, password=ssh_pass)
            print("Get ssh_client by ssh pass.", flush=True)
        else:
            print("Can not connect host without ssh_pass or ssh_private_key!", flush=True)
            sys.exit(1)
    except Exception as err:
        print("连接服务器失败！", flush=True)
        print(err, flush=True)
        exit(1)
    return ssh_client


def checkEnv(sshClient):
    """检查部署主机环境是否符合要求"""
    cmd = driver_check_cmd
    stdin, stdout, stderr = sshClient.exec_command(cmd)
    inspectData = ""
    while True:
        rd = stdout.readline()
        inspectData = inspectData + rd
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(
            "nvidia驱动未安装或状态异常，请在目标主机执行`{0}`检查驱动情况，驱动安装下载请参考https://www.nvidia.cn/Download/index.aspx".format(
                driver_check_cmd), flush=True)
        sys.exit(1)


def container_inspect(sshClient):
    """根据容器名检查是否存在及返回信息"""
    cmd = "{0} inspect {1}".format(docker_path, container_name)
    stdin, stdout, stderr = sshClient.exec_command(cmd)
    inspectData = ""
    while True:
        rd = stdout.readline()
        inspectData = inspectData + rd
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(err, flush=True)
        return None
    return json.loads(inspectData)


def convert_to_byte(size_str):
    """单位换算"""
    units = {
        'b': 1,
        'kb': 1024,
        'k': 1024,
        'mb': 1024 ** 2,
        'm': 1024 ** 2,
        'gb': 1024 ** 3,
        'g': 1024 ** 3,
        'tb': 1024 ** 4,
        't': 1024 ** 4,
        'pb': 1024 ** 5,
        'p': 1024 ** 5,
    }
    size = size_str.lower().strip()
    for unit in units:
        if size.endswith(unit):
            return int(size[:-len(unit)].strip()) * units[unit]
    if size.isdigit():
        return int(size)
    raise ValueError('Invalid size string')


def sameConfig(config):
    """判断当前nemo容器参数与需要部署的是否相同"""
    host_config = config['HostConfig']
    try:
        if jupyter_port != host_config['PortBindings']['8888/tcp'][0]['HostPort']:
            return False
        if int(memlock) != next(filter(lambda ulimit: "memlock" == ulimit['Name'], host_config['Ulimits']), None)[
            'Hard']:
            return False
        if int(stack) != next(filter(lambda ulimit: "stack" == ulimit['Name'], host_config['Ulimits']), None)['Hard']:
            return False
        if convert_to_byte(shm_size) != host_config['ShmSize']:
            return False
    except Exception as e:
        print(e, flush=True)
        return False
    if nemo_image != config['Config']['Image']:
        return False

    return True


def delete_old_container(sshClient, containerName):
    """删除旧nemo容器"""
    print("删除旧容器: {0}".format(containerName), flush=True)
    cmd_deleteContainer = "{0} stop {1} && {0} rm {1}".format(docker_path, containerName)
    print(cmd_deleteContainer, flush=True)
    stdin, stdout, stderr = sshClient.exec_command(cmd_deleteContainer)
    while True:
        rd = stdout.readline()
        print(rd, end="", flush=True)
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(err, flush=True)


def container_deploy(sshClient):
    """部署nemo容器"""
    print("创建容器: {0}".format(container_name), flush=True)
    deploy_cmd = "{0} run --gpus all -d --runtime=nvidia --shm-size={1} -p {2}:8888 --ulimit memlock={3} --ulimit stack={4}  --name={5}  {6} ./start-jupyter.sh".format(
        docker_path, shm_size, jupyter_port, memlock, stack, container_name, nemo_image)
    deploy_cmd = "{0} pull {1} && ".format(docker_path, nemo_image) + deploy_cmd
    print(deploy_cmd, flush=True)
    stdin, stdout, stderr = sshClient.exec_command(deploy_cmd)
    while True:
        rd = stdout.readline()
        print(rd, end="", flush=True)
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print(err, flush=True)
        if "Unknown runtime" in err:
            print(
                "请确保docker版本>19.03并已安装nvidia容器运行时，安装步骤参考https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#setting-up-nvidia-container-toolkit",
                flush=True)
        sys.exit(1)


def check_inspect(sshClient, inspectInfo):
    """检查nemo容器配置"""
    if not inspectInfo or len(inspectInfo) < 1:
        return True
    needRebuild = False
    containerInfo = inspectInfo[0]
    if containerInfo['State']['Status'] != "running" or not containerInfo['State']['Running']:
        needRebuild = True
        print("NeMo container is not running.", flush=True)
    elif not sameConfig(containerInfo):
        needRebuild = True
        print("NeMo config is changed.", flush=True)
    if needRebuild:
        delete_old_container(sshClient, containerInfo['Name'])
        return True
    print("已存在相同实例,不执行任何动作", flush=True)
    return False


def get_container_logs(sshClient, lines=100):
    """获取容器日志"""
    cmd_get_log = f"docker logs {container_name} --tail {lines} 2>&1"
    stdin, stdout, stderr = sshClient.exec_command(cmd_get_log)
    logs = ""
    while True:
        rd = stdout.readline()
        logs = logs + rd
        if not rd:
            break
    err = stderr.read().decode("utf-8")
    if err:
        print("获取容器日志失败！", flush=True)
        print(err, flush=True)
    return logs


def check_inspect_health(sshClient, timeout_seconds):
    start_time = time.time()
    while True:
        if time.time() - start_time >= timeout_seconds:
            print("检测超时！服务启动失败，请前往服务器查看详情！", flush=True)
            logs = get_container_logs(sshClient)
            print(logs, flush=True)
            return False
        try:
            inspectInfo = container_inspect(sshClient)
            if not inspectInfo or len(inspectInfo) < 1:
                print("Container {0} does not exist!".format(container_name), flush=True)
            elif inspectInfo[0]['State']['Health']['Status'] != "healthy":
                print("Health: {0}".format(inspectInfo[0]['State']['Health']['Status']), flush=True)
                logs = get_container_logs(sshClient)
                print(logs, flush=True)
            elif inspectInfo[0]['State']['Health']['Status'] == "healthy":
                print("Health: {0}".format(inspectInfo[0]['State']['Health']['Status']), flush=True)
                return True
        except Exception as e:
            print(e, flush=True)
        time.sleep(30)


if __name__ == '__main__':
    sshClient = getSSHClient()
    checkEnv(sshClient)
    containerInfo = container_inspect(sshClient)
    if check_inspect(sshClient, containerInfo):
        container_deploy(sshClient)
    check_health = check_inspect_health(sshClient, timeout_seconds)
    if check_health:
        print("服务已成功启动！请查看容器日志获取token访问jupyter。", flush=True)
    else:
        sys.exit(1)
